'''
@author: morales
Updated Jan 13, 2010
'''

from dna.Polyspect_Operator import Polyspect
from dna.Parameters import *
import infiniteSitesTools as ist
from dna.PolymorphismTables import *
from dna.PowerTools import Power
from scipy import random
from scipy.linalg import norm
import numpy.matlib as np
from time import gmtime, strftime

class PolyspectSpace():
	def __init__(self, iParameters, iAltDemography, iSFS, iAltSFS ):
		
		self.myPowerPath = []
		self.myParameters = iParameters
		self.myAltDemography = iAltDemography
		self.mySFS = iSFS
		self.myAltSFS = iAltSFS
		self.maxPointPower = None

		n = self.myParameters.myPopulation.numChromosomes
		
		TD = ist.getOp_TajimasD( n )
		self.addPoint('Tajs D', TD )
				
		FWH = ist.getOp_FayWuH( n )
		self.addPoint('FaWu H', FWH )

		FLD = ist.getOp_FuLiD( n )
		self.addPoint('FuLi D', FLD )

		FLF = ist.getOp_FuLiF( n )
		self.addPoint('FuLi F', FLF )
		
		self.myPowerPath.append( self.maxPointPower )
		
		print "Max\t" + str( self.maxPointPower ) + "\n"

#		n = self.myParameters.myPopulation.numOfChromosomes
#		myPolyspect = Polyspect( point, self.myParameters, self.numSimulations, self.mySFS )
#		pointPower = Power( myPolyspect, self.myAltDemography, self.myAltSFS )
#		self.myPowerPath.append( pointPower )

	def addPoint(self, title, linOp):

		point = self.normalize( linOp )
		pointPower = self.getNewPowerPoint( point )
		self.myPowerPath.append( pointPower )

		print str(title) + '\t' + str( pointPower )

		if self.maxPointPower == None:
			self.maxPointPower = pointPower
			
		elif pointPower.myPower >self.maxPointPower.myPower:
			self.maxPointPower = pointPower

	def getCurrentPoint(self):
		point = self.myPowerPath[-1]
		return point
		
	def normalize(self, point):
		
		#ratio = 1.0 / point[1]
		#normPoint = point * ratio
		
		n = norm(point)
		normPoint = point / n
		return normPoint

	def getNextPoint(self):
		
		penultPoint = self.myPowerPath[-2].myPolyspect.myOperator
		lastPoint = self.myPowerPath[-1].myPolyspect.myOperator
		highestPower = self.myPowerPath[-1].myPower
		n = len(lastPoint)

		point = np.ravel(np.zeros(n))
		found = False
		searchCount = 0
		level = 1
		r = 0.001
		
		while not found:
			
			if searchCount == 0:
				print "\nin getNextPoint" + strftime("%a, %d %b %Y %H:%M:%S +0700", gmtime())
				for i in range(1, n):
					point[i] = 2 * lastPoint[i] - penultPoint[i]
			else:
				for i in range(1, n):
					point[i] = lastPoint[i] + r * (random.random() - 0.5)

			normPoint = self.normalize( point )
			pointPower = self.getNewPowerPoint( normPoint )

			if (pointPower.myPower - highestPower) > 0.01:
				self.myPowerPath.append( pointPower )
				found = True
				return pointPower
			
			searchCount += 1
			if searchCount > 100:
				if level < 5:
					level += 1
					print "\nStarting level " + str(level) + " random walk"
					r *= 10.0
					searchCount = 0
				else:
					break

		return None

	def getNewPowerPoint(self, point ):
		n = self.myParameters.myPopulation.numChromosomes
		myPolyspect = Polyspect( point, self.myParameters, self.mySFS )
		pointPower = Power( myPolyspect, self.myAltSFS )
		
		return pointPower
	
	def getSlopeAtMaxPoint(self):
		maxPoint = self.myPowerPath[-1].myPolyspect.myOperator
		n = maxPoint.size
		
		slope = np.ravel(np.zeros(n))
		
		for i in range( 1, n ):
			point1 = maxPoint.copy()
			point1[i] = maxPoint[i] + 0.01
			point1Power = self.getNewPowerPoint( point1 )
			
			point2 = maxPoint.copy()
			point2[i] = maxPoint[i] - 0.01
			point2Power = self.getNewPowerPoint( point2 )
			
			dy = point1Power.myPower - point2Power.myPower
			slope[i] = dy / 0.02
		
		print slope
		return slope
		
	
	def __next__(self):
		if len(self.myPowerPath) < 10:
			return self
		else:
			raise StopIteration
	
	def __iter__(self):
		return self
	
	def __str__(self):
		
		s = ''
		for p in self.myPowerPath:
			s += '\n' + str(p)
		
		return s
